#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6
import time
import torch
from torch.backends import cudnn

from datasets import Data
from nodes import GlobalNode, Node
from options import args_parser
from utils import print_message, Recorder, lr_scheduler
import numpy as np
import random
from fed_avg import train


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.deterministic = True


if __name__ == '__main__':
    # 实验序号
    setup_seed(2020)
    start = time.time()
    args = args_parser()
    print_message(args)  # 打印配置信息
    device = torch.device(args.device)  # 指定gpu
    # 获取数据
    data = Data(args)  # self.testloader==>全部测试集，self.trainloader==>各节点平分数据集
    # noniid情况下，data.trainloader返回dict存储下标
    # iid情况下，data.trainloader返回split后的5*1w样本
    # init nodes
    global_node = GlobalNode(args)
    node_list = []
    for i in range(args.node_num):
        node_list.append(Node(i, data.train_loader[i], data.train_set, args))  # Node0是第一个本地节点，节点的num=1
    # 打印各个节点的模型信息
    for i in range(args.node_num):
        print("Node:{},model:{}".format(i, type(node_list[i].model)))
    print("\nGlobal model:{}".format(type(global_node.model)))

    recorder = Recorder(args)
    loss = []
    acc = []
    for rounds in range(args.R):
        print('===============The {:d}-th round==============='.format(rounds + 1))
        lr_scheduler(rounds, node_list, args)
        # 对各个本地节点训练
        for i in range(len(node_list)):
            node_list[i].fork(global_node)  # 模型下发，每个本地节点copy一份全局模型
            recorder.test_local(node_list[i], data.test_loader)
            recorder.test_meme(node_list[i], data.test_loader)
            recorder.validate_local(node_list[i])  # 本地节点查看本地效果
            recorder.validate_meme(node_list[i])  # meme查看本地效果

            for epoch in range(args.E):  # 本地蒸馏次数
                train(node_list[i], recorder)

        global_node.average(node_list)  # 模型聚合
        # 此时一个epoch已经结束
        # 全局节点的测试集acc和loss
        loss_global, acc_global = global_node.test(data.test_loader)
        loss.append(loss_global)  # 保存meme聚合后测试的acc和loss
        acc.append(acc_global)

    # R个epoch训练结束，打印模型测试的结果
    # recorder.print_acc()  # 打印本地节点表现
    print("训练结束！中心节点测试集结果：\n")
    loss_global, acc_global = global_node.test(data.test_loader)

    print_message(args)

    for i in range(args.node_num):
        print("Node:{},model:{}".format(i, type(node_list[i].model)))
    print("Global model:{}".format(type(global_node.model)))
    end = time.time()
    print("运行时长{}".format(end - start))

    # 保存模型
    for i in range(args.node_num):
        torch.save(node_list[i].model.state_dict(), './save/{}/Node{}.pkl'.format(args.num, i))

    torch.save(global_node.model.state_dict(), './save/{}/meme.pkl'.format(args.num))
    # 保存本地节点验证集的acc和loss
    for i in range(args.node_num + 1):
        np.save('./save/{}/train_loss_{}.npy'.format(args.num, i), recorder.train_loss[str(i)])
        np.save('./save/{}/train_meme_loss_{}.npy'.format(args.num, i),
                recorder.train_meme_loss[str(i)])
        np.save('./save/{}/val_loss_{}.npy'.format(args.num, i), recorder.val_loss[str(i)])
        np.save('./save/{}/val_meme_loss_{}.npy'.format(args.num, i), recorder.val_meme_loss[str(i)])
        np.save('./save/{}/test_loss_{}.npy'.format(args.num, i), recorder.test_loss[str(i)])
        np.save('./save/{}/test_meme_loss_{}.npy'.format(args.num, i),
                recorder.test_meme_loss[str(i)])

        np.save('./save/{}/train_acc_{}.npy'.format(args.num, i), recorder.train_acc[str(i)])
        np.save('./save/{}/train_meme_acc_{}.npy'.format(args.num, i),
                recorder.train_meme_acc[str(i)])
        np.save('./save/{}/val_acc_{}.npy'.format(args.num, i), recorder.val_acc[str(i)])
        np.save('./save/{}/val_meme_acc_{}.npy'.format(args.num, i), recorder.val_meme_acc[str(i)])
        np.save('./save/{}/test_acc_{}.npy'.format(args.num, i), recorder.test_acc[str(i)])
        np.save('./save/{}/test_meme_acc_{}.npy'.format(args.num, i), recorder.test_meme_acc[str(i)])

    # 保存结果，进行实验分析
    np.save('./save/{}/acc_R{}_D{}_G{}_L{}_I{}.npy'.format(args.num, args.R, args.dataset,
                                                           args.global_model, args.local_model, args.iid),
            np.array(acc))
    np.save('./save/{}/loss_R{}_D{}_G{}_L{}_I{}.npy'.format(args.num, args.R, args.dataset,
                                                            args.global_model, args.local_model, args.iid),
            np.array(loss))
